#include "mri_tv.h"

using namespace std;
using namespace cv;
/****************************************************************************/
//Perform forward Fourier Transform and mask
Mat fft2(Mat mask, Mat src)
{
	int nRows = src.rows;
	int nCols = src.cols;
	dft(src, src, DFT_SCALE|DFT_COMPLEX_OUTPUT);
	Mat dst(nRows, nRows, CV_64FC2);
	for (int i=0; i<nRows; i++)
	{
		for (int j=0; j<nCols; j++)
		{
			dst.at<double>(i, j) = src.at<double>(i, j) * mask.at<double>(i, j);
		}

	}

	return dst;
}
/****************************************************************************/
//Perform inverse Fourier Transform and mask
Mat ift2(Mat mask, Mat src)
{
	int nRows = src.rows;
	int nCols = src.cols;
	Mat dst(nRows, nRows, CV_64FC2);
	for (int i=0; i<nRows; i++)
	{
		for (int j=0; j<nCols; j++)
		{
			dst.at<double>(i, j) = src.at<double>(i, j) * mask.at<double>(i, j);
		}
	}
	dft(dst, dst, DFT_INVERSE|DFT_COMPLEX_OUTPUT);
	return dst;
}
/****************************************************************************/
double f2norm(Mat mask, Mat T, Mat S)
{
	double norm2;
	Mat result;
	result = fft2(mask, T);
	result -= S;
	norm2 = norm(result, CV_L2)*norm(result, CV_L2);
	return norm2;
}
/****************************************************************************/
double f1norm_tv(Mat src)
{
	int nRows = src.rows;
	int nCols = src.cols;
	//////////////////////////////////////////////////////////////////////////
	Mat dfxImg(nRows, nCols+2, CV_64FC2);
	copyMakeBorder(src, dfxImg, 0, 0, 1, 1, BORDER_REPLICATE);
	Mat dfx(nRows, nCols, CV_64FC2);
	dfx = 1/2*(dfxImg.colRange(2, nCols+2) - dfxImg.colRange(0, nCols));
	//////////////////////////////////////////////////////////////////////////
	Mat dfyImg(nRows+2, nCols, CV_64FC2);
	copyMakeBorder(src, dfyImg, 1, 1, 0, 0, BORDER_REPLICATE);
	Mat dfy(nRows, nCols, CV_64FC2);
	dfy = 1/2*(dfyImg.rowRange(2, nCols+2) - dfyImg.rowRange(0, nCols));
	//////////////////////////////////////////////////////////////////////////
	Scalar_<double> result;
	Mat tmp;
	sqrt(dfx * dfx + dfy*dfy, tmp);
	result = sum(tmp);
	//////////////////////////////////////////////////////////////////////////
	return result.val[0];// + result.val[1];
}
/****************************************************************************/
Mat g2norm(Mat mask, Mat T, Mat S)
{
	Mat result(T.rows, T.cols, CV_64FC2);
	result = fft2(mask, T);
	result = ift2(mask, result);
	result -= S;
	//cout << result.channels() << endl;
	//cout << S.channels() << endl;
	result *= 2;
	//addWeighted(result, 1, S, -1, 0, result);
	//addWeighted(result, 1, result, 1, 0, result);
	return result;
}
/****************************************************************************/
Mat g1norm_tv(Mat src)
{
	double exps = 0.000000000000001;
	int nRows = src.rows;
	int nCols = src.cols;
	//////////////////////////////////////////////////////////////////////////
	Mat x1Img(nRows, nCols+2, CV_64FC2);
	copyMakeBorder(src, x1Img, 0, 0, 1, 1, BORDER_REPLICATE);
	Mat x1(nRows, nCols, CV_64FC2);
	x1 = x1Img.colRange(1, nCols+1) - x1Img.colRange(2, nCols+2);
	//////////////////////////////////////////////////////////////////////////
	Mat y1Img(nRows+2, nCols, CV_64FC2);
	copyMakeBorder(src, y1Img, 1, 1, 0, 0, BORDER_REPLICATE);
	Mat y1(nRows, nCols, CV_64FC2);
	y1 = y1Img.rowRange(1, nRows+1) - y1Img.rowRange(2, nRows+2);
	//////////////////////////////////////////////////////////////////////////
	//Mat x2Img(nRows, nCols+2, CV_64FC2);
	//copyMakeBorder(src, x2Img, 0, 0, 1, 1, BORDER_REPLICATE);
	//Mat x2(nRows, nCols, CV_64FC2);
	//x2 = x2Img.colRange(0, nCols) - x2Img.colRange(1, nCols+1);
	////////////////////////////////////////////////////////////////////////////
	//Mat y2Img(nRows+2, nCols, CV_64FC2);
	//copyMakeBorder(src, y2Img, 1, 1, 0, 0, BORDER_REPLICATE);
	//Mat y2(nRows, nCols, CV_64FC2);
	//y2 = y2Img.rowRange(1, nRows+1) - y2Img.rowRange(2, nRows+2);
	Mat grad1;
	sqrt(x1*x1 + y1*y1 + exps, grad1);

	divide(x1, grad1, grad1);
	return grad1;
}